# -*- coding: utf-8 -*-
"""
Created on Sat Jun  9 12:46:29 2012

@author: nzhao
"""

from traits.api import HasTraits, Instance, DelegatesTo, Bool, Range, Int, Float, Button, Array, Enum, List, File, String
from traitsui.api import View, Group, HGroup, Item, TableEditor, RangeEditor, ArrayEditor, InstanceEditor
from traitsui.table_column import ObjectColumn, ExpressionColumn
from traitsui.table_filter import EvalFilterTemplate, MenuFilterTemplate, RuleFilterTemplate, EvalTableFilter

import numpy as np

import phy.general.MaterialData as ma
from phy.spin.SingleSpin import SingleSpin
from phy.spin.MultiSpin import MultiSpin
from misc.readxyz import readfile
from math_qm.CoordianteSystem import CoordinateSystem


table_editor = TableEditor(
    columns = [ ObjectColumn( name = 'spin_element', width = 0.05 ),
                ObjectColumn( name = 'x',  width = 0.08 ),
                ObjectColumn( name = 'y',  width = 0.08 ),
                ObjectColumn( name = 'z',  width = 0.08 ),
                ObjectColumn( name = 'spin_type',  width = 0.05 ),
                ObjectColumn( name = 'spin_mark',  width = 0.05 ),
                ObjectColumn( name = 'state_I',  width = 0.1 ),
                ObjectColumn( name = 'state_X',  width = 0.1 ),
                ObjectColumn( name = 'state_Y',  width = 0.1 ),
                ObjectColumn( name = 'state_Z',  width = 0.1 ),
                ExpressionColumn(
                    label = 'Distance',
                    width = 0.20,
                    expression = "float('%f' % ( ( object.x**2 + object.y**2 + object.z**2)**0.5)) "  )
                ],
    deletable   = True,
    sort_model  = False,
    auto_size   = False,
    orientation = 'vertical',
    edit_view   = View(
                      Group(HGroup( 'spin_element', 'spin_type', 'x', 'y', 'z', 'spin_mark',\
                              'state_I','state_X','state_Y','state_Z',
                             show_border = True
                      )),
                      resizable = True
                  ),

    filters     = [ EvalFilterTemplate, MenuFilterTemplate, RuleFilterTemplate ],
    search      = EvalTableFilter(),
    show_toolbar = True,
    row_factory = SingleSpin
    )

spin_viewer_group = Group(
            Group(
                HGroup(
                    Item('is_use_file', show_label = False),
                    Item('file_name', show_label = True, enabled_when = 'is_use_file'),
                    
                    ),
                HGroup(
                    Item('is_Lattice', show_label = False), 
                    Item('lattice_const',enabled_when='is_Lattice', label='Use lattice constant'),
                    Item('random_state', label='Random Pure State', show_label = False),
                    Item('thermal_state', label='Thermal Mixed State', show_label = False),
                    Item('list_update', show_label=False),
                    ),
                Item( 'spin_list', show_label  = False, editor = table_editor),
                label='Spin List'
                ),
            Group(
                HGroup( 
                    Item('index1', style='simple', label='Spin-A index:'), 
                    Item('index2', style='simple', label='Spin-B index:') 
                    ),
                HGroup(
                    Item('axes_T', editor = InstanceEditor( name = 'axes_list', editable = False), 
                         label='axes', style = 'custom'),
                    Item('distance', style='readonly', label='Distance between Spin-A and Spin-B [A]= '),
                    Item('tensor', style='readonly', label='  Dipolar coupling between Spin-A and Spin-B = '),
                    Item('tensor_unit')
                    ),
                HGroup(
                    Item('axes_S', editor = InstanceEditor( name = 'axes_list', editable = False), 
                         label='axes', style = 'custom'),
                    Item('S_comp', style='custom', label='The tensor projected on '),
                    ),
                HGroup(
                    Item('axes_hf', editor = InstanceEditor( name = 'axes_list', editable = False), 
                         label='axes', style = 'custom'),
                    Item('hf_norm', style='readonly', label=' effective field vector = '),
                    Item('str_1', style='readonly', show_label = False),
                    Item('hf_dir', style='readonly', show_label = False),
                    Item('str_2', style='readonly', show_label = False),
                    Item('copy_2_custom', show_label=False, label= 'customize')
                    ),
                label='Calculations'
                ),
            show_border = True,
            label = 'Spin Viewer'
        )
        
        
class SpinViewer(HasTraits):
    is_use_file   = Bool(False)    
    file_name     = File('/home/nzhao/research/computation/workspace/cce-engine/input/bath/C12Abund99_99/')

    is_Lattice    = Bool(False)    
    lattice_const = Float(ma.DIAMOND_LATTICE_CONSTANT)
    
    random_state  = Button
    thermal_state = Button
    
    spin_list     = List (SingleSpin)
    multi_spin    = Instance(MultiSpin, ())
    list_update   = Button()
    
    maxN          = Int
    index1        = Range(low=1, high='maxN', editor = RangeEditor(low=1, high_name = 'maxN', mode = 'spinner'))
    index2        = Range(low=1, high='maxN', editor = RangeEditor(low=1, high_name = 'maxN', mode = 'spinner'))
    
    axes_T   = Instance( CoordinateSystem, () )
    axes_S   = Instance( CoordinateSystem, () )
    axes_hf  = Instance( CoordinateSystem, () )
    axes_list  = List( CoordinateSystem )

    
    distance      = Float(0.0)
    tensor        = Array(float,(3,3),np.zeros((3,3)), editor = ArrayEditor(width = 150, format_str = '%8.3f'))
    tensor_unit   = Enum('k rad/s', 'Hz', 'kHz', 'MHz', 'GHz')

    S_comp        = Array(value = np.array([1.0, 1.0, 1.0])/np.sqrt(3.), editor = ArrayEditor(width = 50, format_str = '%8.5f'))    
    hf_dir        = Array(value = np.zeros(3),               editor = ArrayEditor(width = 50, format_str = '%8.5f'))
    hf_norm       = Float
    str_1         = String('* [')
    str_2         = String(']')
    copy_2_custom = Button

    def __init__(self, axes = None):
        if axes is not None:
            self.axes_list = axes
        else:
            self.axes_list = [CoordinateSystem()]
        
        self.axes_T  = self.axes_list[0]
        self.axes_S  = self.axes_list[0]
        self.axes_hf = self.axes_list[0]
        
        self.spin_list   = self.multi_spin.spin_list
        self.update_multispin()
    
    def _copy_2_custom_fired(self):
        self.axes_list[2].x_axis = self.hf_dir
        
    def _random_state_fired(self):
        np.random.seed(1)
        for spin in self.multi_spin.spin_list:
            v=np.random.rand(3)
            spin.state_X, spin.state_Y, spin.state_Z = (v / np.linalg.norm(v))
            spin.state_I = 0.0
    
    def _thermal_state_fired(self):
        for spin in self.multi_spin.spin_list:
            spin.state_I, spin.state_X, spin.state_Y, spin.state_Z = 1.0, 0.0, 0.0, 0.0
        
    def _file_name_changed(self):
        bath_spin_list = readfile(self.file_name)
        self.spin_list = self.spin_list + bath_spin_list
        self.update_multispin()
        
    def _spin_list_items_changed(self):
        self.update_multispin()

    def _list_update_fired(self):
        self.update_multispin()

    def _axes_T_changed(self):
        self.update_calculation()
    def _axes_S_changed(self):
        self.update_calculation()
    def _axes_hf_changed(self):
        self.update_calculation()


    def _index1_changed(self):      self.update_calculation()
    def _index2_changed(self):      self.update_calculation()
    def _tensor_unit_changed(self): self.update_calculation()        
    
    def update_multispin(self):
        self.multi_spin.spin_list = self.spin_list
        
        if self.is_Lattice:
            for spin in self.multi_spin.spin_list:
                spin.x *= self.lattice_const
                spin.y *= self.lattice_const
                spin.z *= self.lattice_const
        
        self.maxN = len(self.spin_list)
        self.multi_spin.spin_num = self.maxN

        spin_type_list = [x.spin_type for x in self.spin_list]
        self.multi_spin.center_spin_index = spin_type_list.index('Central')

        self.update_calculation()
        
    def update_calculation(self):
        spin1         = self.multi_spin.spin_list[self.index1-1]
        spin2         = self.multi_spin.spin_list[self.index2-1]
        self.distance = self.multi_spin.dist(spin1, spin2)
        self.tensor   = self.multi_spin.dipolar(spin1, spin2, self.axes_T)
        
        v = self.axes_S._vector(self.S_comp)
        dir_vector = v / np.linalg.norm(v)
        
        hf_vector = self.axes_T._vector( np.dot(self.axes_T._component(dir_vector), self.tensor ) )
        
        self.hf_norm = np.linalg.norm(hf_vector)
        
        if self.hf_norm > 0.0:
            self.hf_dir  = self.axes_hf._component( hf_vector / self.hf_norm)
        else:
            self.hf_dir = np.array([0., 0., 1.])
        
        if self.tensor_unit == 'kHz':
            self.tensor = self.tensor / (2.0 * np.pi)
            self.hf_norm = self.hf_norm / (2.0 * np.pi)
        elif self.tensor_unit == 'Hz':
            self.tensor = self.tensor / (2.0 * np.pi / 1000.0)
            self.hf_norm = self.hf_norm / (2.0 * np.pi / 1000.0)
        elif self.tensor_unit == 'MHz':
            self.tensor = self.tensor / (2.0 * np.pi * 1000.0)
            self.hf_norm = self.hf_norm / (2.0 * np.pi * 1000.0)
        elif self.tensor_unit == 'GHz':
            self.tensor = self.tensor / (2.0 * np.pi * 10**6)
            self.hf_norm = self.hf_norm / (2.0 * np.pi * 10**6)
    
    view = View(
        spin_viewer_group,
        title     = 'Spin Viewer',
        width     = .8,
        height    = .9,
        resizable = True,
        buttons   = [ 'OK' ],
        kind      = 'live'
    )

if __name__ == '__main__':
    sv=SpinViewer()
    sv.configure_traits()